package at.ltasat.classifier.dt;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import at.ltasat.classifier.Classifier;
import at.ltasat.classifier.Clazz;
import at.ltasat.classifier.ClazzManager;
import at.ltasat.classifier.Membership;
import at.ltasat.classifier.ObjectFeatures;
import at.ltasat.classifier.ObjectFeaturesFactory;
import at.ltasat.classifier.Trainer;

final class DecisionTreeTrainer<O, D> extends Trainer<O, D> {

	private static final int SPLIT_MODE_MIDDLE = 0, SPLIT_MODE_MEAN = 1, SPLIT_MODE_AVERAGE = 3;

	private final Map<ObjectFeatures, Membership<D>[]> data;
	private final int splitMode;

	public DecisionTreeTrainer(String name, ClazzManager<D> classManager, ObjectFeaturesFactory<O> factory) {
		super(name, classManager, factory);
		data = new HashMap<ObjectFeatures, Membership<D>[]>();
		splitMode = SPLIT_MODE_MEAN;
	}

	@Override
	public void train(final ObjectFeatures features, Membership<D>... memberships) {
		assert Membership.check(memberships);
		memberships = memberships.clone();
		if (data.containsKey(features)) {
			final Map<Clazz<D>, Double> degrees = new HashMap<Clazz<D>, Double>();
			for (Membership<D> membership : memberships) {
				degrees.put(membership.getClazz(), membership.getDegree());
			}
			for (Membership<D> membership : data.get(features)) {
				final Clazz<D> clazz = membership.getClazz();
				if (degrees.containsKey(clazz)) {
					degrees.put(clazz, degrees.get(clazz) + membership.getDegree());
				}
			}
			for (Clazz<D> clazz : degrees.keySet()) {
				degrees.put(clazz, degrees.get(clazz) / 2);
			}
		}
		assert Membership.check(memberships);
		Arrays.sort(memberships);
		data.put(features, memberships);
	}

	@Override
	public Classifier<O, D> create() {
		return new DecisionTreeClassifier<O, D>(getName(), getClazzManager(), getObjectFeaturesFactory(), learn());
	}

	private Node<D> learn() {
		final List<ObjectFeatures> examples = new ArrayList<ObjectFeatures>();
		final boolean[] attributes = new boolean[getNumberOfDimensions()];
		examples.addAll(data.keySet());
		Arrays.fill(attributes, true);
		return learn(examples, attributes);
	}

	private Node<D> learn(final List<ObjectFeatures> examples, final boolean[] attributes) {
		Node<D> node;
		final Membership<D>[] commonClassification = getCommonClassification(examples, false);
		if (commonClassification != null) {
			node = new ClassificationNode<D>(commonClassification);
		} else if (!isAnyAttributeLeft(attributes)) {
			node = new ClassificationNode<D>(getCommonClassification(examples, true));
		} else {
			final int feature = chooseFeature(examples, attributes);
			final double split = chooseSplit(examples, feature);
			final List<ObjectFeatures> lowerExamples = new ArrayList<ObjectFeatures>(), upperExamples = new ArrayList<ObjectFeatures>();
			splitExamples(feature, split, examples, lowerExamples, upperExamples);
			attributes[feature] = false;
			final Node<D> lower = learn(lowerExamples, attributes);
			final Node<D> upper = learn(upperExamples, attributes);
			attributes[feature] = true;
			node = new DecisionNode<D>(feature, split, lower, upper);
		}
		return node;
	}

	@SuppressWarnings("unchecked")
	private Membership<D>[] getDefault() {
		final double[] degrees = new double[getClazzManager().size()];
		for(ObjectFeatures feature : data.keySet()){
			for(Membership<D> membership : data.get(feature)){
				degrees[membership.getClazz().getId()] += membership.getDegree();
			}
		}
		for(Clazz<D> clazz : getClazzManager()){
			degrees[clazz.getId()] = degrees[clazz.getId()] / data.size();
		}
		final Membership<D>[] memberships = (Membership<D>[]) new Membership<?>[getClazzManager().size()];
		for (Clazz<D> clazz : getClazzManager()) {
			memberships[clazz.getId()] = new Membership<D>(clazz, degrees[clazz.getId()]);
		}
		return memberships;
	}

	@SuppressWarnings("unchecked")
	private Membership<D>[] getCommonClassification(final List<ObjectFeatures> examples, final boolean allowContradictions) {
		if (examples.size() == 0) {
			return getDefault();
		}
		final ObjectFeatures representative = examples.get(0);
		final Clazz<D>[] clazzes = getClazzes(representative);
		final double[] degrees = getDegrees(representative);
		for (int i = 1; i < examples.size(); i++) {
			final Membership<D>[] memberships = data.get(examples.get(i));
			final int nClazzes = Math.min(clazzes.length, memberships.length);
			for (int clazzIndex = 0; clazzIndex < nClazzes; clazzIndex++) {
				final Clazz<D> expectedClazz = clazzes[clazzIndex];
				final Membership<D> membership = memberships[clazzIndex];
				if (!allowContradictions && expectedClazz != membership.getClazz()) {
					return null;
				}
				degrees[clazzIndex] += membership.getDegree();
			}
		}
		for (int i = 0; i < degrees.length; i++) {
			degrees[i] = degrees[i] / examples.size();
		}
		final Membership<D>[] common = (Membership<D>[]) new Membership<?>[clazzes.length];
		for (int i = 0; i < clazzes.length; i++) {
			common[i] = new Membership<D>(clazzes[i], degrees[i]);
		}
		return common;
	}

	@SuppressWarnings("unchecked")
	private Clazz<D>[] getClazzes(final ObjectFeatures example) {
		final Membership<D>[] memberships = data.get(example);
		final Clazz<D>[] clazzes = (Clazz<D>[]) new Clazz<?>[memberships.length];
		for (int i = 0; i < clazzes.length; i++) {
			clazzes[i] = memberships[i].getClazz();
		}
		return clazzes;
	}

	private double[] getDegrees(final ObjectFeatures example) {
		final Membership<D>[] memberships = data.get(example);
		final double[] degrees = new double[memberships.length];
		for (int i = 0; i < degrees.length; i++) {
			degrees[i] = memberships[i].getDegree();
		}
		return degrees;
	}

	private boolean isAnyAttributeLeft(final boolean[] attributes) {
		for (boolean attribute : attributes) {
			if (attribute) {
				return true;
			}
		}
		return false;
	}

	private int chooseFeature(final List<ObjectFeatures> examples, final boolean[] attributes) {
		int feature = 0;
		while (!attributes[feature]) {
			feature++;
		}
		return feature;
	}

	private double chooseSplit(final List<ObjectFeatures> examples, final int feature) {
		Collections.sort(examples, new FeatureComparator(feature));
		final double split;
		switch (splitMode) {
		case SPLIT_MODE_MIDDLE:
			split = getMiddle(feature, examples, 0, examples.size() - 1);
			break;
		case SPLIT_MODE_AVERAGE:
			double sum = 0;
			for (ObjectFeatures example : examples) {
				sum += example.getFeatureValue(feature);
			}
			split = sum / examples.size();
			break;
		case SPLIT_MODE_MEAN:
			split = getMiddle(feature, examples, examples.size() / 2, (examples.size() + 1) / 2);
			break;
		default:
			assert false;
			return 0;
		}
		return split;
	}

	private static double getMiddle(int feature, List<ObjectFeatures> examples, int index1, int index2) {
		final double lower = examples.get(index1).getFeatureValue(feature);
		final double upper = examples.get(index2).getFeatureValue(feature);
		assert lower <= upper;
		return lower + (upper - lower) / 2;
	}

	private void splitExamples(final int feature, final double split, final List<ObjectFeatures> examples, final List<ObjectFeatures> lowerExamples, final List<ObjectFeatures> upperExamples) {
		for (ObjectFeatures example : examples) {
			if (example.getFeatureValue(feature) < split) {
				lowerExamples.add(example);
			} else {
				upperExamples.add(example);
			}
		}
	}

	private static final class FeatureComparator implements Comparator<ObjectFeatures> {

		private final int feature;

		public FeatureComparator(int feature) {
			this.feature = feature;
		}

		@Override
		public int compare(ObjectFeatures f1, ObjectFeatures f2) {
			return (int) Math.signum(f1.getFeatureValue(feature) - f2.getFeatureValue(feature));
		}

	}

}
